import pandas as pd
import pickle
# decision tree
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import classification_report

df = pd.read_csv('server/test.csv', encoding='gbk')
print(df.shape)
target = df['职业']
feature = pd.get_dummies(df.loc[:,['年龄','学历','婚姻状况','性别','出生地']])
pkl_filename = "server/model.pkl"

# one-hot encoding
print(feature.head())
print(feature.columns)
X = feature.values
y = target.values
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

clf = DecisionTreeClassifier()
clf.fit(X_train, y_train)
# save model

y_pred = clf.predict(X_test)
print(classification_report(y_test, y_pred))
with open(pkl_filename, 'wb') as file:
    pickle.dump(clf, file)

with open(pkl_filename, 'rb') as file:
    pickle_model = pickle.load(file)
# Calculate the accuracy score and predict target values
y_pred = pickle_model.predict(X_test)
print(classification_report(y_test, y_pred))